{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}

module FRP.Taupo.Awaiting
  ( AwaitingSignal,
    signal,
    signalMono,
    signalSemi,
    signalEarliest,
    signalLatest,
    event,
    output,
  )
where

import Control.Applicative (liftA2)
import Control.Foldl (Fold (Fold))
import qualified Control.Foldl as Foldl
import Data.Bifunctor (bimap)
import Data.Functor (void)
import Data.Functor.Apply (Apply ((<.>)))
import Data.Functor.Combinator (HFunctor (hmap), Inject (inject))
import Data.Kind (Type)
import Data.Monoid (Alt (Alt), Dual (Dual), getAlt, getDual)
import Data.Semialign (Repeat (repeat), Semialign (align, alignWith), Zip (zip, zipWith))
import Data.These (These (That, These, This))
import FRP.Taupo.Async (AsyncSignal)
import Prelude hiding (init, repeat, zipWith)

data AwaitingSignal (eff :: Type -> Type) (a :: Type) where
  AwaitingPure :: eff a -> AwaitingSignal eff a
  AwaitingSignal :: AsyncSignal eff a -> Fold a b -> AwaitingSignal eff b
  AwaitingEvent :: eff b -> AwaitingSignal eff a -> AwaitingSignal eff a
  AwaitingMux ::
    (a -> b -> c) ->
    AwaitingSignal eff a ->
    AwaitingSignal eff b ->
    AwaitingSignal eff c

instance (Functor eff) => Functor (AwaitingSignal eff) where
  fmap f = \case
    AwaitingPure act -> AwaitingPure (fmap f act)
    AwaitingSignal sig fld -> AwaitingSignal sig (fmap f fld)
    AwaitingEvent act sig -> AwaitingEvent act (fmap f sig)
    AwaitingMux g sig1 sig2 -> AwaitingMux (\x -> f . g x) sig1 sig2

instance (Apply eff) => Semialign (AwaitingSignal eff) where
  align sig1 sig2 = case sig1 of
    AwaitingPure act -> case sig2 of
      AwaitingPure act' -> AwaitingPure (These <$> act <.> act')
      _ -> AwaitingMux These sig1 sig2
    AwaitingSignal sig fld -> case sig2 of
      AwaitingSignal sig' fld' ->
        AwaitingSignal (align sig sig') (uncurry These <$> asyncMash fld fld')
      _ -> AwaitingMux These sig1 sig2
    _ -> AwaitingMux These sig1 sig2
  alignWith f sig1 sig2 = case sig1 of
    AwaitingPure act -> case sig2 of
      AwaitingPure act' -> AwaitingPure ((\x -> f . These x) <$> act <.> act')
      _ -> AwaitingMux (\x -> f . These x) sig1 sig2
    AwaitingSignal sig fld -> case sig2 of
      AwaitingSignal sig' fld' ->
        AwaitingSignal
          (align sig sig')
          (f . uncurry These <$> asyncMash fld fld')
      _ -> AwaitingMux (\x -> f . These x) sig1 sig2
    _ -> AwaitingMux (\x -> f . These x) sig1 sig2

instance (Apply eff) => Zip (AwaitingSignal eff) where
  zip sig1 sig2 = case sig1 of
    AwaitingPure act -> case sig2 of
      AwaitingPure act' -> AwaitingPure ((,) <$> act <.> act')
      _ -> AwaitingMux (,) sig1 sig2
    AwaitingSignal sig fld -> case sig2 of
      AwaitingSignal sig' fld' -> AwaitingSignal (align sig sig') (asyncMash fld fld')
      _ -> AwaitingMux (,) sig1 sig2
    _ -> AwaitingMux (,) sig1 sig2
  zipWith f sig1 sig2 = case sig1 of
    AwaitingPure act -> case sig2 of
      AwaitingPure act' -> AwaitingPure (f <$> act <.> act')
      _ -> AwaitingMux f sig1 sig2
    AwaitingSignal sig fld -> case sig2 of
      AwaitingSignal sig' fld' ->
        AwaitingSignal (align sig sig') (uncurry f <$> asyncMash fld fld')
      _ -> AwaitingMux f sig1 sig2
    _ -> AwaitingMux f sig1 sig2

instance (Applicative eff, Apply eff) => Repeat (AwaitingSignal eff) where
  repeat = AwaitingPure . pure

instance (Applicative eff, Apply eff) => Applicative (AwaitingSignal eff) where
  pure = repeat
  liftA2 = zipWith
  (<*>) = zipWith ($)

instance HFunctor AwaitingSignal where
  hmap nt = \case
    AwaitingPure act -> AwaitingPure (nt act)
    AwaitingSignal sig fld -> AwaitingSignal (hmap nt sig) fld
    AwaitingEvent act sig -> AwaitingEvent (nt act) (hmap nt sig)
    AwaitingMux f sig1 sig2 -> AwaitingMux f (hmap nt sig1) (hmap nt sig2)

instance Inject AwaitingSignal where
  inject = AwaitingPure

signal :: AsyncSignal eff a -> Fold a b -> AwaitingSignal eff b
signal = AwaitingSignal

signalMono :: (Monoid m) => (a -> m) -> AsyncSignal eff a -> AwaitingSignal eff m
signalMono f sig = AwaitingSignal sig (Foldl.foldMap f id)

signalSemi ::
  (Semigroup s) =>
  (a -> s) ->
  AsyncSignal eff a ->
  AwaitingSignal eff (Maybe s)
signalSemi f sig = AwaitingSignal sig (Foldl.foldMap (Just . f) id)

signalEarliest :: AsyncSignal eff a -> AwaitingSignal eff (Maybe a)
signalEarliest sig = AwaitingSignal sig (Foldl.foldMap (Alt . Just) getAlt)

signalLatest :: AsyncSignal eff a -> AwaitingSignal eff (Maybe a)
signalLatest sig = AwaitingSignal sig (Foldl.foldMap (Dual . Alt . Just) (getAlt . getDual))

event :: eff () -> AwaitingSignal eff a -> AwaitingSignal eff a
event = AwaitingEvent

output :: (Functor eff) => AwaitingSignal eff a -> AwaitingSignal eff ()
output = void

-- Helpers

asyncMash :: forall a b c d. Fold a b -> Fold c d -> Fold (These a c) (b, d)
asyncMash (Fold step init extract) (Fold step' init' extract') =
  Fold
    ( \(l, r) -> \case
        This y -> (step l y, r)
        That z -> (l, step' r z)
        These y z -> (step l y, step' r z)
    )
    (init, init')
    (bimap extract extract')
